# Copyright (c) 2023 Presto Labs Pte. Ltd.
# Author: leon

import datetime
import logging
import signal
from typing import List

import tornado.httpserver
import tornado.websocket
import tornado.ioloop
import tornado.web
from absl import app, flags

from coin.exchange.uniswap_base.rest.subscriber import UniswapSubscriber
from coin.exchange.uniswap_base.constants import PRODUCTION_RPC_URL
from coin.flow.subscriber import from_queue
from coin.flow.simple_queue import SimpleQueue

FLAGS = flags.FLAGS


class WsHandler(tornado.websocket.WebSocketHandler):
  clients = []
  def open(self):
    logging.info('New connection.')
    self.clients.append(self)

  def on_message(self, message):
    logging.info('Message received: %s' % message)

  def on_close(self):
    logging.info('Connection closed.')
    self.clients.remove(self)

  @classmethod
  def write_to_clients(cls, data):
    for client in cls.clients:
      client.write_message(data)


class RawFeedHandler(object):
  _clients: List[WsHandler]

  def __init__(self, clients: List[WsHandler]):
    self._clients = clients

  def on_book(self, record, queue_data, topic_data):
    for client in self._clients:
      client.write_message(record.data)

  def on_trade(self, record, queue_data, topic_data):
    for client in self._clients:
      client.write_message(record.data)


def main(argv):
  queue = SimpleQueue('uniswap-v2-publisher')
  ioloop = tornado.ioloop.IOLoop.current()
  logger = logging.getLogger('uniswapv2_publisher')

  worker_id = '1'
  sub = UniswapSubscriber(
      queue,
      'a',
      worker_id,
      FLAGS.symbols,
      provider=PRODUCTION_RPC_URL['ethereum'],
      close_callback=None,
      logger=logger)
  sub.start(ioloop)

  flow_sub = from_queue(queue)
  raw_feed_handler = RawFeedHandler(WsHandler.clients)
  flow_sub.subscribe('uniswapv2_rest_depth_*%s' % worker_id, raw_feed_handler.on_book)
  flow_sub.subscribe('uniswapv2_rest_trade_*%s' % worker_id, raw_feed_handler.on_trade)

  def sig_handler(self, sig=None, frame=None):
    ioloop.add_callback_from_signal(ioloop.stop)

  ioloop.add_timeout(datetime.timedelta(minutes=FLAGS.exit_after_min), ioloop.stop)
  signal.signal(signal.SIGTERM, sig_handler)
  signal.signal(signal.SIGINT, sig_handler)

  application = tornado.web.Application([
    (r'/ws', WsHandler),
  ])
  http_server = tornado.httpserver.HTTPServer(application)
  http_server.listen(9443)

  ioloop.start()


if __name__ == '__main__':
  flags.DEFINE_integer('exit_after_min', 1440, 'exit_after_min')

  flags.DEFINE_list('symbols', 'CCX-USDT', "Separate by comma.")

  logging.basicConfig(level='INFO', format='%(levelname)s %(asctime)s %(name)s] %(message)s')
  app.run(main)
